/**
  SUNNY project, Anyang Normal University, IMP-CAS
  \class TASparseVec
  \brief Serving as a column for sparse matrix, inherited from ROOT main I/O class
   - TTree, which is dedicated to handle large amount of data. Note that one object
  stands for one column in a sparse matrix, so each entry is a 3-member tuple. Zero
  elements are just skipped.
  \author SUN Yazhou, asia.rabbit@163.com
  \since 2021/02/10
  \date Last modified: 2022/12/04 by SUN Yazhou
  \copyright 2020-2021 SUN Yazhou
  \copyright SUNNY project, Anyang Normal University, IMP-CAS
*/

#include <algorithm>
#include <string>
#include <cmath>
#include <iomanip>
#include "config.h" // contains user switch of definition of ROOT
#ifdef CATCH_CNOK
#include "catch.h"
#endif
#include "TASparseVec.h"
#include "TASparseMatrix.h"
#include "TAMath.h"
#include "TAException.h"

#ifdef _ROOT_
#include "TTree.h"
#include "TSQLResult.h"
#include "TSQLRow.h"
#endif

using std::cout;
using std::endl;
using std::setw;
using std::ios_base;
using std::stoul;

TASparseVec::TASparseVec(const string &name, const string &title)
    : TATreeCol(name, title){
  CreateBranch();
} // end constructor

TASparseVec::TASparseVec(const TASparseVec &t) : TATreeCol(t){
  CreateBranch();
  CopyFill(&t); // fill t to this
} // end of copy constructor

// the move constructor
TASparseVec::TASparseVec(TASparseVec &&v) : TATreeCol(std::move(v)){
  SetBranchAddress();
} // end move ctor

// copy v to this
TASparseVec &TASparseVec::operator=(const TASparseVec &v){
  TATreeCol::operator=(v);

  CreateBranch();
  CopyFill(&v); // fill v to this
  return *this;
} // end assignment function


void TASparseVec::CreateBranch(){
#ifdef _ROOT_
  fTree->Branch("index", &fIndex, "index/l");
  fTree->Branch("r", &fR, "r/l");
  fTree->Branch("v", &fV, "v/D");
#else
  TAException::Error("TASparseVec", "CreateBranch: ROOT not found.");
#endif
} // end member function CreateBranch

void TASparseVec::SetBranchAddress(){
#ifdef _ROOT_
  // already set
  if((void*)(&fV) == (void*)fTree->GetBranch("v")->GetAddress()) return;

  fTree->SetBranchAddress("index", &fIndex);
  fTree->SetBranchAddress("r", &fR);
  fTree->SetBranchAddress("v", &fV);
#else
    TAException::Error("TASparseVec", "SetBranchAddress: ROOT not found.");
#endif

  // for(ull i = 0; i < GetEntries(); i++) GetEntry(i);
  // if(GetEntries() > 0) GetEntry(0); // seal up the connection
} // end member function SetBranchAddress

/// \retval row of the last (non-zero) entry
ull TASparseVec::GetMaxRow() const{
  return GetRow(GetEntries()-1);
} // end member function GetMaxRow

// trim v to zero if fabs(v) < epsilon (e.g. 1.e-16)
double TASparseVec::trim(double v){
  if(fabs(v) < 1.e-15) return 0.;
  return v;
} // end member funtion trim

void TASparseVec::Fill(ull r, double v){
  v = trim(v);
  if(!v) return;

  fIndex = GetEntries();
  fR = r; fV = v;
  SetBranchAddress();
  TATreeCol::Fill();
} // end member function Fill

// get the value of index (row) r
double TASparseVec::GetValue(ull r) const{
#ifdef _ROOT_
  char selection[512];
  sprintf(selection, "r==%llu", r);
  TSQLResult *res = fTree->Query("index", selection);
  const int nr = res->GetRowCount();
  if(nr > 1) TAException::Warn("TASparseVec", // found multiple entries
    "GetValue: multiple (%d) entries found for selection %s", nr, selection);
  if(!nr) return 0.; // found none - perceived as an unregistered zero value

  GetEntry(std::stoul(res->Next()->GetField(0))); delete res;
#else
    TAException::Error("TASparseVec", "GetValue: ROOT not found.");
#endif
  return fV;
} // end member function GetValue

/// \retval fR of entryIndex
ull TASparseVec::GetRow(ull entryIndex) const{
  if(entryIndex >= GetEntries()){
    TAException::Warn("TASparseVec", "GetRow: Required entry index out of border.");
    return 0;
  } // end if

  GetEntry(entryIndex);
  return fR;
} // end member function GetRow

// \retval this += k*v
void TASparseVec::SelfAdd(double k, const TASparseVec &v){
#ifdef _ROOT_
  // obtain the entry index range with fV != 0. or v.fV != 0. //
  if(!v.GetEntries() || !k) return;
  if(!GetEntries()) v.Scale(k, *this); // r = 0, and r += k*v => r = k*v
  const ull n1 = std::min(GetRow(0), v.GetRow(0));
  const ull n2 = std::max(GetMaxRow(), v.GetMaxRow());

  TASparseVec t;
  for(ull i = n1; i <= n2; i++) t.Fill(i, GetValue(i) + k*v[i]);

  // transfer ownership of fTree
  delete fTree; fTree = t.fTree; t.fTree = nullptr;
  SetBranchAddress();
  // fTree->Scan(); Print(); // DEBUG
#else
  TAException::Error("TASparseVec", "SelfAdd: ROOT not found.");
#endif
} // end member function SelfAdd

// r = k*this
void TASparseVec::Scale(double k, TASparseVec &r) const{
  if(!GetEntries()) return; // this is a zero vector, so just leave r empty is enough
  r.Clear(); // remove all the entries in r
  const ull n = GetEntries();
  for(ull i = 0; i < n; i++){
    GetEntry(i);
    r.Fill(fR, k*fV);
  } // end for over i
} // end member function Scale

/// \retval (*this)*v
double TASparseVec::DotProduct(const TASparseVec &v) const{
  if(!GetEntries()) return 0.;
  const ull n1 = std::max(GetRow(0), v.GetRow(0));
  const ull n2 = std::min(GetMaxRow(), v.GetMaxRow());
  double r = 0., e1, e2;
  for(ull i = n1; i <= n2; i++)
    if((e1 = GetValue(i)) && (e2 = v[i])) r += e1 * e2;
  return r;
} // end member function DotProduct

void TASparseVec::CopyFill(const TATreeCol *t){
  Clear();
  if(!t->GetEntries()) return;

  const TASparseVec *tt = dynamic_cast<const TASparseVec *>(t);
  if(!tt) TAException::Error("TASparseVec", "CopyFill: Not a TASparseVec pointer");

  const ull n1 = tt->GetRow(0), n2 = tt->GetMaxRow();
  double e;
  for(ull i = n1; i <= n2; i++) if((e = (*tt)[i])) Fill(i, e);
} // end member function CopyFill


// \retval sqrt(norm2)
double TASparseVec::norm() const{
  return sqrt(norm2());
} // end member function norm

// \retval (*this)*(*this)
double TASparseVec::norm2() const{
  if(!GetEntries()) return 0.;
  const ull n1 = GetRow(0);
  const ull n2 = GetMaxRow();
  double r = 0.;
  for(ull i = n1; i <= n2; i++) r += TAMath::sqr(GetValue(i));
  return r;
} // end member function norm2

// remove q bits from r: r-= |q><q|r>
void TASparseVec::Purify(const TASparseVec &q){
  // zero vector is orthogonal to any vector, need no reorthogonalizations
  if(!GetEntries()) return;
  SelfAdd(-DotProduct(q), q);
} // end member function Purify

// remove Q=(q1,q2..qn) bits from r: r-= \sum_q{|q><q|r>}
void TASparseVec::Purify(const TASparseMatrix &Q){
  // zero vector is orthogonal to any vector, need no reorthogonalizations
  if(!GetEntries()) return;
  for(TASparseVec *q : Q) Purify(*q);
} // end member function Purify

void TASparseVec::Print() const{
  // cout << "Branch status - index: " << fTree->GetBranchStatus("index") << endl;
  // cout << "Branch status - r: " << fTree->GetBranchStatus("r") << endl;
  // cout << "Branch status - v: " << fTree->GetBranchStatus("v") << endl;
  const int nn = GetEntries(), NM = 20;
  const int n = NM < nn ? NM : nn;
  cout << "TASparseVec Print: totally " << nn << " elements." << endl;
  if(n > NM) cout << "(Display only the first " << NM << " entries.)" << endl;
  ios_base::fmtflags initial = cout.setf(ios_base::fixed, ios_base::floatfield);
  // cout.unsetf(ios_base::floatfield);
  cout.precision(6);
  cout << std::right;
  for(int i = 0; i < n; i++){
    GetEntry(i);
    cout << setw(10) << fR;
  }
  cout << "\033[32;1m" << endl;
  for(int i = 0; i < n; i++){
    GetEntry(i);
    cout << setw(10) << fV;
  }
  cout << "\033[0m" << endl;
  cout.setf(initial);
  // fTree->Print();
  // fTree->Scan();
} // end member function Print

// print in C/C++ form
void TASparseVec::PrintInC(ull maxLen) const{
  const ull n = GetEntries();
  ios_base::fmtflags initial = cout.setf(ios_base::fixed, ios_base::floatfield);
  cout.unsetf(ios_base::floatfield);
  cout.precision(6);
  ull cnt = 0; // Nof nonzero elements printed
  cout << "{";
  for(ull i = 0; i < n; i++){
    GetEntry(i);
    while(cnt < fR){ cout << "0., "; cnt++; } // skipped elements are all zeros
    if(i < n-1) cout << fV << ", ";
    else cout << fV;
    cnt++;
  } // end for over i
  cout << "};" << endl;
  cout.setf(initial);
} // end member function PrintInC

void TASparseVec::normalize(){
#ifdef _ROOT_
  if(!GetEntries()) return;
  const double p = norm();
  if(1. == p) return;
  if(0. == p) return; // zero vector - need not to normalize

  TASparseVec tmp;
  Scale(1./p, tmp);
  delete fTree;
  fTree = tmp.fTree;
  tmp.fTree = nullptr;
  SetBranchAddress();
#else
  TAException::Error("TASparseVec", "SelfAdd: ROOT not found.");
#endif
} // end member function normalize


#ifdef _ROOT_
#ifdef CATCH_CNOK
// unit tests with -catch-
TEST_CASE("sparse vector tests", "[spv]"){
  TASparseVec a;
  a.Fill(0, 3.);
  a.Fill(1, 4.);
  a.Fill(2, 1.e-20);
  CHECK(2 == a.GetEntries());
  CHECK(0. == a[2]);
  CHECK(a.norm() == Approx(5.).epsilon(1.e-15));
  // scale and selfAdd
  TASparseVec r;
  a.Scale(2., r);
  CHECK(r[0] == 6.);
  CHECK(a[1] == 4.);
  CHECK(a.DotProduct(r) == 50.);
  a.SelfAdd(2., r);
  r.normalize(); a.Purify(r);
  CHECK(double(a[0]+1.) == Approx(1.).epsilon(1.e-8));
  CHECK(double(a[1]+1.) == Approx(1.).margin(1.e-8));
} // end TEST_CASE
#endif
#endif
